// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/exec/group_id.h>

namespace kumo::pollux::exec {
    GroupId::GroupId(
        int32_t operatorId,
        DriverCtx *driverCtx,
        const std::shared_ptr<const core::GroupIdNode> &groupIdNode)
        : Operator(
            driverCtx,
            groupIdNode->outputType(),
            operatorId,
            groupIdNode->id(),
            "GroupId") {
        const auto &inputType = groupIdNode->sources()[0]->outputType();

        melon::F14FastMap<column_index_t, column_index_t>
                outputToInputGroupingKeyMapping;
        for (const auto &groupingKeyInfo: groupIdNode->groupingKeyInfos()) {
            outputToInputGroupingKeyMapping[outputType_->getChildIdx(
                        groupingKeyInfo.output)] =
                    inputType->getChildIdx(groupingKeyInfo.input->name());
        }

        auto numGroupingKeys = groupIdNode->numGroupingKeys();

        groupingKeyMappings_.reserve(groupIdNode->groupingSets().size());
        for (const auto &groupingSet: groupIdNode->groupingSets()) {
            std::vector<column_index_t> mappings(numGroupingKeys, kMissingGroupingKey);
            for (const auto &groupingKey: groupingSet) {
                auto outputChannel = outputType_->getChildIdx(groupingKey);
                POLLUX_USER_CHECK_NE(
                    outputToInputGroupingKeyMapping.count(outputChannel),
                    0,
                    "GroupIdNode didn't map grouping key {} to input channel",
                    groupingKey);
                auto inputChannel = outputToInputGroupingKeyMapping.at(outputChannel);
                mappings[outputChannel] = inputChannel;
            }

            groupingKeyMappings_.emplace_back(std::move(mappings));
        }

        const auto &aggregationInputs = groupIdNode->aggregationInputs();
        aggregationInputs_.reserve(aggregationInputs.size());
        for (auto i = 0; i < aggregationInputs.size(); ++i) {
            const auto &input = aggregationInputs[i];
            aggregationInputs_.push_back(inputType->getChildIdx(input->name()));
        }
    }

    bool GroupId::needsInput() const {
        return !noMoreInput_ && input_ == nullptr;
    }

    void GroupId::addInput(RowVectorPtr input) {
        // Load Lazy vectors.
        for (auto &child: input->children()) {
            child->loaded_vector();
        }

        input_ = std::move(input);
    }

    RowVectorPtr GroupId::getOutput() {
        if (!input_) {
            return nullptr;
        }

        // Make a copy of input for the grouping set at 'groupingSetIndex_'.
        auto numInput = input_->size();

        std::vector<VectorPtr> outputColumns(outputType_->size());

        const auto &mapping = groupingKeyMappings_[groupingSetIndex_];
        auto numGroupingKeys = mapping.size();

        // Fill in grouping keys.
        for (auto i = 0; i < numGroupingKeys; ++i) {
            if (mapping[i] == kMissingGroupingKey) {
                // Add null column.
                outputColumns[i] = BaseVector::create_null_constant(
                    outputType_->childAt(i), numInput, pool());
            } else {
                outputColumns[i] = input_->childAt(mapping[i]);
            }
        }

        // Fill in aggregation inputs.
        for (auto i = 0; i < aggregationInputs_.size(); ++i) {
            outputColumns[numGroupingKeys + i] = input_->childAt(aggregationInputs_[i]);
        }

        // Add groupId column.
        outputColumns[outputType_->size() - 1] =
                std::make_shared<ConstantVector<int64_t> >(
                    pool(), numInput, false, BIGINT(), groupingSetIndex_);

        ++groupingSetIndex_;
        if (groupingSetIndex_ == groupingKeyMappings_.size()) {
            groupingSetIndex_ = 0;
            input_ = nullptr;
        }

        return std::make_shared<RowVector>(
            pool(), outputType_, nullptr, numInput, std::move(outputColumns));
    }
} // namespace kumo::pollux::exec
